#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：JobShopEnv.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/8 20:20 
'''
import numpy as np
import random

class JobShopEnv:
    def __init__(self, num_jobs, num_machines):
        self.num_jobs = num_jobs  # 作业数量
        self.num_machines = num_machines  # 机器数量
        self.makespan = 0  # 总完成时间
        self.jobs = self._create_jobs()
        self.machines = [None] * num_machines  # 每个机器当前处理的工件

    def _create_jobs(self):
        # 为每个作业随机生成工艺路线和处理时间
        jobs = []
        for _ in range(self.num_jobs):
            # 每个工件有一个随机的加工工艺，包含每道工序的机器和处理时间
            job = [(random.randint(0, self.num_machines - 1), random.randint(1, 10)) for _ in range(self.num_machines)]
            jobs.append(job)
        return jobs

    def reset(self):
        # 重置环境
        self.makespan = 0
        self.machines = [None] * self.num_machines
        self.jobs = self._create_jobs()
        return self._get_state()

    def _get_state(self):
        # 当前状态可以是每个机器的工件状态
        state = [machine if machine is None else machine[0] for machine in self.machines]
        return np.array(state)

    def step(self, action):
        # 根据动作更新环境状态
        job_idx, machine_idx = action
        if self.machines[machine_idx] is None:  # 如果机器空闲
            self.machines[machine_idx] = self.jobs[job_idx].pop(0)  # 分配工件

        done = all(len(job) == 0 for job in self.jobs)  # 所有作业完成则结束
        reward = -1  # 简单的负奖励模型，鼓励尽快完成任务

        if done:
            self.makespan = self._calculate_makespan()

        return self._get_state(), reward, done, {}

    def _calculate_makespan(self):
        # 计算总完成时间（Makespan）
        return sum([task[1] for machine in self.machines if machine is not None for task in machine])
